{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell first to import some modules\n",
    "# we'll need to run the examples below.\n",
    "%matplotlib inline\n",
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import HTML\n",
    "\n",
    "from pydrake.all import (AddMultibodyPlantSceneGraph,\n",
    "                         DiagramBuilder,\n",
    "                         Parser,\n",
    "                         Simulator,\n",
    "                         SignalLogger)\n",
    "from underactuated import ( FindResource, PlanarSceneGraphVisualizer )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This cell sets up a double pendulum, runs a simulation, and\n",
    "# renders a video of its results.\n",
    "\n",
    "# Set up a block diagram with the robot (dynamics) and a visualization block.\n",
    "builder = DiagramBuilder()\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder)\n",
    "\n",
    "# Load the double pendulum from Universal Robot Description Format\n",
    "parser = Parser(plant, scene_graph)\n",
    "parser.AddModelFromFile(FindResource(\"double_pendulum/double_pendulum.urdf\"))\n",
    "plant.Finalize()\n",
    "\n",
    "builder.ExportInput(plant.get_actuation_input_port())\n",
    "visualizer = builder.AddSystem(PlanarSceneGraphVisualizer(scene_graph,\n",
    "                                                          xlim=[-2.8, 2.8],\n",
    "                                                          ylim=[-2.8, 2.8]))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                visualizer.get_input_port(0))\n",
    "\n",
    "# The logger needs to be told to expect a 4-element input\n",
    "# (the 4-element double pendulum state, in this case).\n",
    "logger = builder.AddSystem(SignalLogger(4))\n",
    "logger.DeclarePeriodicPublish(0.033333, 0.0)\n",
    "builder.Connect(plant.get_state_output_port(), logger.get_input_port(0))\n",
    "\n",
    "diagram = builder.Build()\n",
    "\n",
    "# Set up a simulator to run this diagram\n",
    "simulator = Simulator(diagram)\n",
    "\n",
    "# Set the initial conditions\n",
    "context = simulator.get_mutable_context()\n",
    "context.SetContinuousState([1., 1., 0., 0.])  # (theta1, theta2, theta1dot, theta2dot)\n",
    "context.FixInputPort(0, [0.,0.])   # Zero input torques\n",
    "\n",
    "# Simulate and animate\n",
    "visualizer.start_recording(show=False)\n",
    "simulator.AdvanceTo(10.0)\n",
    "ani = visualizer.get_recording(repeat=True)\n",
    "HTML(ani.to_html5_video())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell after running a simulation to plot traces\n",
    "# of the pendulum state across time.\n",
    "fields = [\"shoulder\", \"elbow\"]\n",
    "for i in range(2):\n",
    "    #plt.subplot(3, 1, i+1)\n",
    "    plt.plot(logger.sample_times(), logger.data()[(i, i+2), :].transpose())\n",
    "    plt.legend([\"position\", \"velocity\"])\n",
    "    plt.xlabel('t')\n",
    "    plt.ylabel(fields[i])\n",
    "    plt.grid(True)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
